use fastmcp_rs::{
    FastMcpServer,
    run_stdio,
    tool::ToolResponse,
    mcp_register_tools, mcp_prompt_text, mcp_register_prompts, mcp_response, mcp_text,
};
use serde_json::{Value, json};
use tokio::signal;

// 将工具抽离至模块顶层，便于复用与测试
#[fastmcp_rs_macros::mcp_tool(
    name = "add",
    description = "Adds two numbers together",
    parameters_json(
        "type": "object",
        "required": ["a", "b"],
        "properties": {
            "a": { "type": "number" },
            "b": { "type": "number" }
        }
    ),
    annotations(category = json!("math"))
)]
async fn add(_: fastmcp_rs::InvocationContext, payload: Value) -> fastmcp_rs::Result<ToolResponse> {
    let a = payload.get("a").and_then(Value::as_f64).ok_or_else(|| {
        fastmcp_rs::FastMcpError::InvalidInvocation(
            "expected float argument 'a'".into(),
        )
    })?;
    let b = payload.get("b").and_then(Value::as_f64).ok_or_else(|| {
        fastmcp_rs::FastMcpError::InvalidInvocation(
            "expected float argument 'b'".into(),
        )
    })?;
    let sum = a + b;
    Ok(mcp_response!([ mcp_text!(format!("{a} + {b} = {sum}")) ]))
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .init();

    let server = FastMcpServer::builder()
        .name("RustAdd")
        .instructions("Call the add tool to compute the sum of two numbers.")
        .build()
        .into_shared();

    // 通过宏生成的 `add_tool()` 完成注册
    // 若启用 `auto-register` 特性，工具会在构建服务器时自动注册，无需重复注册
    #[cfg(not(feature = "auto-register"))]
    {
        mcp_register_tools!(server, [ add_tool() ]);
    }

    mcp_register_prompts!(server, [
        mcp_prompt_text! {
            name: "addition",
            role: "user",
            text: "Please add {{ a }} and {{ b }}.",
            description: "Asks the assistant to add two numbers",
            parameters: json!({
                "type": "object",
                "required": ["a", "b"],
                "properties": {
                    "a": { "type": "number" },
                    "b": { "type": "number" }
                }
            }),
        }
    ]);

    tokio::select! {
        result = run_stdio(server.clone()) => result?,
        _ = signal::ctrl_c() => {
            tracing::info!("received Ctrl+C, shutting down");
        }
    }

    Ok(())
}
